import asyncio
import os
import time

import aiofiles
import aiohttp
import requests
from Crypto.Cipher import AES


def get_m3u8(url):
    res = requests.get(url)
    with open("../数据暂存/m3u8.txt", "wb") as f:
        f.write(res.content)
    print("m3u8文件下载完成")


def get_key():
    #  key : [49, 66, -84, -94, 86, 35, -77, 74, 96, -128, 89, 42, -79, -94, 26, 98]
    with open("../数据暂存/m3u8.txt", "r", encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line.startswith('#EXT-X-KEY'):
                iv = line.strip('IV=0x')[1]
                iv = bytes.fromhex(iv)
                break
    key = b''
    for i in b:
        i = i.to_bytes(1, 'big', signed=True)
        key += i
    return key, iv


def dec_ts(key, iv, data):
    """AES解密视频"""
    aes = AES.new(key, IV=iv, mode=AES.MODE_CBC)
    data = aes.decrypt(data)
    return data


async def down_ts(url, session, key):
    """异步下载单个ts文件"""
    name = url[-6:]
    n = 1
    while n == 1:
        try:
            async with session.get(url, headers=head) as res:
                if res.status != 200:
                    print("响应异常", res.status)
                    raise
                async with aiofiles.open(f"../数据暂存/{name}", "wb") as f:
                    data = await res.content.read()
                    data = dec_ts(*key, data)
                    await f.write(data)
            print(f"{name}下载完成")
            # 记录已下载的文件名称
            global lst
            lst.append(name)
            n = 0
        except Exception as e:
            print(f"{name}下载出现异常", e)
            time.sleep(3)


async def start_down_ts(key, url):
    up_url = url.rsplit('/', maxsplit=1)[0]
    tasks = []
    # 避免出现 ServerDisconnected
    # 降低并发数量为50
    connector = aiohttp.TCPConnector(limit=50)
    # 设置超时时间为600s
    # timeout = aiohttp.ClientTimeout(total=600)
    async with aiohttp.ClientSession(connector=connector) as sess:

        async with aiofiles.open("../数据暂存/m3u8.txt", "r", encoding="utf-8") as f:
            line: str
            n = 0
            async for line in f:
                if line.startswith("#"):
                    continue
                    n += 1
                ts_url = up_url + '/' + line.strip()
                # print(ts_url)
                task = asyncio.create_task(down_ts(ts_url, sess, n, key))
                tasks.append(task)
            await asyncio.wait(tasks)


def merge(name):
    """合并ts文件，生成完整的mp4文件"""
    # 获取所有ts文件列表
    global lst
    lst.sort()
    s = "+".join(lst)
    os.chdir("../数据暂存")
    os.system(f"copy /b {s} {name}")
    os.system(f"move {name} E:\T\Videos\一人之下")
    print("合并成功")
    # 删除暂存数据
    os.system("del *.ts")
    os.system("del *.txt")
    print("======Done=====")


def main(url, name):
    # 获取m3u8
    get_m3u8(url)
    # 获取密钥
    key = get_key()
    # 下载源文件
    loop = asyncio.get_event_loop()
    loop.run_until_complete(start_down_ts(key, url))
    # 视频整合
    merge(name)


if __name__ == '__main__':
    head = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/108.0.0.0 Safari/537.36"}
    url = 'https://cd11-ccd1-2.play.bokecc.com/flvs/44CACEB725BAB224/2022-06-16/3241AFA15520B04963835A29B2A11961-90.m3u8?t=1673102019&key=B394B6EB7C37656B823078041B29EBC2&tpl=10&tpt=112&custom_id=Z_8985cfec064011edbdcb506b4bfed66c_204993'
    # 定义一个列表，记录ts文件名称
    lst = []
    t1 = time.time()
    main(url, "接口测试.mp4")
    t2 = time.time()
    print(f"总耗时{t2 - t1}s")
